from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv1D, LSTM, Dense, Dropout, GlobalMaxPooling1D
from tensorflow.keras.wrappers.scikit_learn import KerasClassifier
from tensorflow.keras.optimizers import Adam
from sklearn.model_selection import StratifiedKFold

class DeepLearningTrainer:
    def __init__(self, input_shape, model_type='cnn'):
        self.input_shape = input_shape
        self.model_type = model_type
    
    def build_model(self):
        """构建深度学习模型"""
        if self.model_type == 'cnn':
            model = Sequential([
                Conv1D(64, 3, activation='relu', input_shape=self.input_shape),
                Dropout(0.5),
                Conv1D(128, 3, activation='relu'),
                GlobalMaxPooling1D(),
                Dense(128, activation='relu'),
                Dense(1, activation='sigmoid')
            ])
        elif self.model_type == 'rnn':
            model = Sequential([
                LSTM(128, return_sequences=True, input_shape=self.input_shape),
                Dropout(0.3),
                LSTM(64),
                Dense(64, activation='relu'),
                Dense(1, activation='sigmoid')
            ])
        model.compile(
            optimizer=Adam(0.001),
            loss='binary_crossentropy',
            metrics=['accuracy']
        )
        return model
    
    def cross_validate(self, X, y, epochs=20, batch_size=32, n_splits=3):
        """深度学习交叉验证"""
        kfold = StratifiedKFold(n_splits=n_splits)
        scores = []
        
        for train_idx, test_idx in kfold.split(X, y):
            model = self.build_model()
            history = model.fit(
                X[train_idx], y[train_idx],
                epochs=epochs,
                batch_size=batch_size,
                validation_data=(X[test_idx], y[test_idx]),
                verbose=0
            )
            scores.append(history.history['val_accuracy'][-1])
            
        return {'accuracy': np.mean(scores)}
    
    def save_model(self, model, path):
        """保存Keras模型"""
        model.save(path) 